#!/usr/bin/env python

# Imports:
# -------
import getopt, sys, os, time, subprocess, re
from xml.etree import ElementTree 

# ----------------------------
def parseCommandLineArgs(argv):
# ----------------------------
	""" Parse the array of command line arguments for the input file """	

	# Parse the command line argument(s)
	filePath = None
	try:
		opts, args = getopt.getopt(argv[1:], 'f:', ['file'])
	except getopt.GetoptError as e:
		print '[ERROR] - %s' % str(e)
		return None

	for opt, val in opts:
		if opt in ['-f', '--file']: filePath = val

	if filePath is None: 
		print '[ERROR] - Must specify a file'
		return None

	return filePath

# --------------------------
def parseInputFile(filePath):
# --------------------------
	""" Parse the input file for the bags, launchfiles, and options """

	# Initialize the containers we're reading the XML file for
	bags = []
	options = {}

	# Load the XML file
	if os.path.exists(filePath):
		# Pull out the root
		elementTree = ElementTree.parse(filePath)
		root = elementTree.getroot()

		# Parse
		for element in root:
			tag = element.tag

			# Bags:
			# ----
			if tag == 'bags':
				for e in element:
					# Add the bag
					if e.tag == 'bag':
						name = e.attrib['name']
						launchFiles = [ee.text for ee in e]
						bags.append((name, launchFiles))
					
					else: 
						print '[ERROR] - Unrecognized tag: %s' % e.tag
						return None

			elif tag == 'options':
				for e in element:
					if e.tag == 'launch_pause': 
						options[e.tag] = float(e.text)
					else: 
						options[e.tag] = e.text
						

		bagBasePath = '' if 'bag_base_path' not in options else options['bag_base_path']
		launchBasePath = '' if 'launch_base_path' not in options else options['launch_base_path']
		if 'initial_launch' in options: options['initial_launch'] = os.path.join(launchBasePath, options['initial_launch'])

		# Combine the basepath and the bag files
		bags = [(os.path.join(bagBasePath, bag[0]), [os.path.abspath(os.path.join(launchBasePath, b)) for b in bag[1]]) for bag in bags]
		return  (bags, options)
	
	else: 
		print '[ERROR] - Specified file "%s" could not be found' % filePath
		return None

# --------------------------------
def killLaunchfile(launchFilePath):
# --------------------------------
	# Load the launch file & parse out the node names
	elementTree = ElementTree.parse(launchFilePath)
	root = elementTree.getroot()
	nodes = root.findall('node')

	nodeNames = []
	for node in nodes:
		nodeNamespace = '' if 'ns' not in node.attrib else (node.attrib['ns'] + '/')
		nodeName = nodeNamespace  + node.attrib['name']
		nodeNames.append(nodeName)

	# Use `rosnode kill` to kill the nodes
	for nodeName in nodeNames:
		subprocess.call('rosnode kill %s' % nodeName, shell=True, stdout = subprocess.PIPE)

# -------------------
def parseOutput(text):
# -------------------
	# Search the output for a sequence of lines of the form ##, ##, ##
	matches = re.findall('([0-9]*\.?[0-9]+), ([0-9]*\.?[0-9]+), ([0-9]*\.?[0-9]+ *)', text)

	# Format the matches
	matches = [[float(m) for m in match] for match in matches]
	return matches

# -------------------------
def writeOutput(path, data):
# -------------------------
	# Open the file for writing
	outFile = open(path, 'w')
	
	for row in data:
		outFile.write('%f, %f, %f\n' % (row[0], row[1], row[2]))

	outFile.close()


# ------------
def main(argv):
# ------------
	# Print the header
	print '               _        ____                  _                          _         '
	print '    /\        | |      |  _ \                | |                        | |        '
	print '   /  \  _   _| |_ ___ | |_) | ___ _ __   ___| |__  _ __ ___   __ _ _ __| | __     '
	print '  / /\ \| | | | __/ _ \|  _ < / _ \ \'_ \ / __| \'_ \| \'_ ` _ \ / _` | \'__| |/ / '
	print ' / ____ \ |_| | || (_) | |_) |  __/ | | | (__| | | | | | | | | (_| | |  |   <      '
	print '/_/    \_\__,_|\__\___/|____/ \___|_| |_|\___|_| |_|_| |_| |_|\__,_|_|  |_|\_\\    '
	print '																				      '                                                                         
	
	# Parse the command line arguments & load the argument file
	filePath = parseCommandLineArgs(argv)
	bags, options = parseInputFile(filePath)
	outDir = '' if 'output_directory' not in options else options['output_directory']

	# If we have an initial launch, launch it
	iLaunchProcess = None
	if 'initial_launch' in options:
		print ' + Initial launchfile: %s' % options['initial_launch']
		launchCommand = 'roslaunch %s' % options['initial_launch']
		iLaunchProcess = subprocess.Popen(launchCommand, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

	# Loop through the bag files and run the launchfile for each
	for i, bagTup in enumerate(bags):
		# Split the bagTup into its components
		bag, launchList = bagTup
		print ' + Bag (%i/%i): %s' % (i + 1, len(bags), bag)
		
		if os.path.exists(bag):
			# Loop through the launchfiles 
			for j, launchfile in enumerate(launchList):
				print '   - Launchfile (%i/%i): %s' % (j + 1, len(launchList), launchfile)

				# Make sure the launchfile exists
				if os.path.exists(launchfile):			
					# Run the launchfile
					launchCommand = 'roslaunch %s' % launchfile
					print '    * %s' % launchCommand
					launchProcess = subprocess.Popen(launchCommand, shell = True, stdout=subprocess.PIPE, bufsize=-1, stderr=subprocess.PIPE)

					# Play the bag file
					bagCommand = 'rosbag play %s' % bag
					print '    * %s' % bagCommand
					bagProcess = subprocess.Popen(bagCommand, shell = True, stdout = subprocess.PIPE, bufsize=-1, stderr = subprocess.PIPE)

					# Wait for the bag process to end
					bagOut, bagErr = bagProcess.communicate()
					
					# Kill the launchfile
					killLaunchfile(launchfile)

					# Get the contents of stdout/stderr from the launchfile
					launchOut, launchErr = launchProcess.communicate()
					launchOut = parseOutput(launchOut)

					# Write the (cleaned up) output to disk
					bagPath, bagName = os.path.split(bag)
					bagName, bagExt = os.path.splitext(bagName)
					launchPath, launchName = os.path.split(launchfile)
					launchName, launchExt = os.path.splitext(launchName)
					outputFilePath = os.path.join(outDir, bagName + '-' + launchName)
					writeOutput(outputFilePath, launchOut)
							
				else: 
					print '     * [ERROR] - Could not find launchfile: "%s"' % launchfile
		else: 
			print '     * [ERROR] - Could not find bag file: "%s"' % bag

	
	if iLaunchProcess != None:
		print ' + Ending initial launchfile: %s' % options['initial_launch']
		killLaunchfile(options['initial_launch'])
	

# -----------------------
if __name__ == '__main__':
# -----------------------
	sys.exit(main(sys.argv))
